import os
import shutil
import time

from flask import send_file
from minio import Minio

from airflow.configuration import conf
from airflow.models.file import File
from airflow.utils.log.logging_mixin import LoggingMixin


class S3Resource(LoggingMixin):

    def __init__(self):
        super().__init__()
        self.s3_client = None
        self.s3_bucket = None
        self.s3_folder = None
        self.resources_folder = conf.get('core', 'resources_folder')

    def get_s3_client(self):
        if self.s3_client is None:
            resources_minio_folder = conf.get('core', 'resources_minio_folder')
            if resources_minio_folder and resources_minio_folder.startswith('s3://'):
                self.s3_client = Minio(conf.get('minio', 'endpoint'),
                                       access_key=conf.get('minio', 'access_key'),
                                       secret_key=conf.get('minio', 'secret_key'),
                                       secure=False)
                prefix = resources_minio_folder.split('//')[1]
                if not prefix.endswith('/'):
                    prefix += '/'
                if not self.resources_folder.endswith('/'):
                    self.resources_folder += '/'
                self.s3_bucket = prefix.split('/')[0]
                self.s3_folder = prefix[prefix.find('/') + 1:]

        return self.s3_client

    def download(self, source: str, target_folder: str):
        if source.startswith('/'):
            source = source[1:]
        if not target_folder.endswith('/'):
            target_folder += '/'
        try:
            s3_client = self.get_s3_client()
            if s3_client is None:
                source = self.resources_folder + source
                self.copy(source, target_folder + source.split('/')[-1])
            else:
                if source.startswith('s3://'):
                    pre = source[5:]
                    bucket = pre.split('/')[0]
                    key = pre[pre.find('/') + 1:]
                    s3_client.fget_object(bucket, key, target_folder + source.split('/')[-1])
                else:
                    source = self.s3_folder + source
                    self.log.info('download file from ' + ' s3://{}/{}'.format(self.s3_bucket, source))
                    # fix back fill in case download file at the same time
                    try:
                        s3_client.fget_object(self.s3_bucket, source, target_folder + source.split('/')[-1])
                    except Exception as ex:
                        time.sleep(1)
                        if os.path.exists(target_folder + source.split('/')[-1]):
                            self.log.info('file already exist {}'.format(source))
                        else:
                            raise ex
        except Exception as e:
            self.log.exception(e)

    def upload(self, source: str, target: str):
        if target.startswith('/'):
            target = target[1:]
        try:
            s3_client = self.get_s3_client()
            if s3_client is None:
                target = self.resources_folder + target
                self.copy(source, target)
            else:
                if target.startswith('s3://'):
                    pre = target[5:]
                    bucket = pre.split('/')[0]
                    key = pre[pre.find('/') + 1:]
                    s3_client.fput_object(bucket, key, source)
                else:
                    target = self.s3_folder + target
                    self.log.info('upload file to ' + ' s3://{}/{}'.format(self.s3_bucket, target))
                    s3_client.fput_object(self.s3_bucket, target, source)
        except Exception as e:
            self.log.exception(e)

    def copy(self, source, target):
        self.log.info('get resource from {}'.format(source))
        os.system('cp {} {}'.format(source, target))

    # 查询file
    def list_file(self, path: str):
        s3_client = self.get_s3_client()
        if s3_client is None:
            return self.os_list_file(path)
        else:
            return self.s3_list_file(path)

    def s3_list_file(self, path: str):
        s3_client = self.get_s3_client()
        if (path.strip() == ''):
            prefix = self.s3_folder
        else:
            prefix = self.s3_folder + path
        objects = s3_client.list_objects(self.s3_bucket, prefix)
        files = []
        for ob in objects:
            if (ob.is_dir):
                obj_name = ob.object_name.split('/')[-2] + "/"
                is_dir = True
            else:
                obj_name = ob.object_name.split('/')[-1]
                is_dir = False
            last_modified = ""
            if (ob.last_modified != None):
                last_modified = ob.last_modified.strftime("%Y-%m-%d %H:%M:%S")
            project = ob.object_name.split('/')[1]

            file = File(obj_name, is_dir, ob.size, last_modified, path + obj_name, project)
            files.append(file)
        mapFile = map(File.to_json, files)
        return list(mapFile)

    def os_list_file(self, path: str):
        if (path.strip() == ''):
            prefix = self.resources_folder
        else:
            prefix = self.resources_folder + path
        dirs_info = os.scandir(prefix + path)
        files = []
        for ob in dirs_info:
            if (os.path.isdir(ob.path)):
                obj_name = ob.name + os.sep
                is_dir = True
            else:
                obj_name = ob.name
                is_dir = False
            last_modified = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(os.path.getctime(ob.path)))
            project = path.split(os.sep)[0]
            size = os.path.getsize(ob.path)
            file = File(obj_name, is_dir, size, last_modified, path + obj_name, project)
            files.append(file)
        mapFile = map(File.to_json, files)
        return list(mapFile)

    # 删除文件
    def delete_file(self, path: str):
        s3_client = self.get_s3_client()
        if s3_client is None:
            self.os_delete_file(path)
        else:
            self.s3_delete_file(path)

    def os_delete_file(self, path: str):
        if (path.strip() == ''):
            prefix = self.resources_folder
        else:
            prefix = self.resources_folder + path
        delete_path = prefix + path
        if (os.path.exists(delete_path)):
            shutil.rmtree(delete_path)

    def s3_delete_file(self, path: str):
        s3_client = self.get_s3_client()
        if (path.strip() == ''):
            prefix = self.s3_folder
        else:
            prefix = self.s3_folder + path
        self.s3_remove_objects(prefix, s3_client)

    def s3_remove_objects(self, path: str, s3_client):
        if (path.endswith(".") or path.endswith("/")):
            objects = s3_client.list_objects(self.s3_bucket, path)
            for ob in objects:
                if (ob.is_dir):
                    self.s3_remove_objects(ob.object_name, s3_client)
                else:
                    s3_client.remove_object(self.s3_bucket, ob.object_name)
        else:
            s3_client.remove_object(self.s3_bucket, path)

    # 上传文件
    def upload_file(self, target: str, file, size: int):
        if target.startswith('/'):
            target = target[1:]
        s3_client = self.get_s3_client()
        if s3_client is None:
            target = self.resources_folder + target
            file.save(os.path.join(target, file.filename))
        else:
            if target.startswith('s3://'):
                pre = target[5:]
                bucket = pre.split('/')[0]
                key = pre[pre.find('/') + 1:] + file.filename
                s3_client.put_object(bucket, key, file, size)
            else:
                target = self.s3_folder + target + file.filename
                self.log.info('upload file to ' + ' s3://{}/{}'.format(self.s3_bucket, target))
                s3_client.put_object(self.s3_bucket, target, file, size)

    # 上传文件
    def download_file(self, source: str):
        if source.startswith('/'):
            source = source[1:]
        s3_client = self.get_s3_client()
        if s3_client is None:
            source = self.resources_folder + source
            if (os.path.exists(source)):
                return send_file(source, as_attachment=True)
        else:
            if source.startswith('s3://'):
                pre = source[5:]
                bucket = pre.split('/')[0]
                key = pre[pre.find('/') + 1:]
                ob = s3_client.get_object(bucket, key)
            else:
                source = self.s3_folder + source
                self.log.info('upload file to ' + ' s3://{}/{}'.format(self.s3_bucket, source))
                ob = s3_client.get_object(self.s3_bucket, source)
                obj_name = source.split('/')[-1]
            return send_file(ob, as_attachment=True, attachment_filename=obj_name)


if __name__ == '__main__':
    s3c = S3Resource()
    s3c.download("flink.jar", "D:\\data")
